use std::collections::HashMap;

use syn::{
    BinOp, Expr, ExprBinary, ExprIf, ExprLit, ExprParen, ExprPath, ExprReturn, FnArg, ItemFn, Lit,
    Pat, PatIdent, ReturnType, Stmt, Type,
};

use crate::mir::{
    BinOpKind, CmpOpKind, MirFunction, MirParam, MirType, Operand, RValue, RegisterAllocator,
    Statement,
};

pub fn generate_mir(source: &str) -> Result<String, String> {
    let file = syn::parse_file(source).map_err(|err| err.to_string())?;

    let func = file
        .items
        .into_iter()
        .find_map(|item| match item {
            syn::Item::Fn(item_fn) => Some(item_fn),
            _ => None,
        })
        .ok_or_else(|| "只支持包含单个函数的源码文件".to_string())?;

    let mut builder = FunctionBuilder::new(&func)?;
    builder.lower_function_body(&func)?;
    let mir = builder.finish()?;
    Ok(mir.format())
}

struct FunctionBuilder {
    mir: MirFunction,
    variables: HashMap<String, u32>,
    reg_alloc: RegisterAllocator,
    current_block: usize,
    return_block: Option<usize>,
    ret_ty: MirType,
}

impl FunctionBuilder {
    fn new(func: &ItemFn) -> Result<Self, String> {
        let ret_ty = match &func.sig.output {
            ReturnType::Default => MirType::Unit,
            ReturnType::Type(_, ty) => map_type(ty)?,
        };

        let mut mir = MirFunction::new(func.sig.ident.to_string(), ret_ty);
        let mut variables = HashMap::new();

        let mut next_reg = 1u32;
        for input in func.sig.inputs.iter() {
            let (name, ty) = match input {
                FnArg::Typed(pat_type) => {
                    let ident = match &*pat_type.pat {
                        Pat::Ident(PatIdent { ident, .. }) => ident.to_string(),
                        _ => return Err("不支持的参数模式".to_string()),
                    };
                    let ty = map_type(&pat_type.ty)?;
                    (ident, ty)
                }
                FnArg::Receiver(_) => {
                    return Err("不支持方法形式的 self 参数".to_string());
                }
            };
            let reg = next_reg;
            next_reg += 1;
            mir.params.push(MirParam {
                name: name.clone(),
                ty,
                register: reg,
            });
            variables.insert(name, reg);
        }

        let entry = mir.add_block();
        debug_assert_eq!(entry, 0, "入口基本块索引应为0");

        Ok(Self {
            mir,
            variables,
            reg_alloc: RegisterAllocator::new(next_reg),
            current_block: entry,
            return_block: None,
            ret_ty,
        })
    }

    fn finish(mut self) -> Result<MirFunction, String> {
        if let Some(return_idx) = self.return_block {
            if let Some(block) = self.mir.blocks.get_mut(return_idx) {
                if !matches!(block.statements.last(), Some(Statement::Return)) {
                    block.statements.push(Statement::Return);
                }
            }
        }

        if self.ret_ty == MirType::I32 && !self.has_explicit_return() {
            return Err("缺少返回语句".to_string());
        }
        Ok(self.mir)
    }

    fn has_explicit_return(&self) -> bool {
        self.mir
            .blocks
            .iter()
            .any(|block| matches!(block.statements.last(), Some(Statement::Return)))
    }

    fn lower_function_body(&mut self, func: &ItemFn) -> Result<(), String> {
        self.lower_block(&func.block, false)
    }

    fn lower_block(&mut self, block: &syn::Block, in_branch: bool) -> Result<(), String> {
        for (index, stmt) in block.stmts.iter().enumerate() {
            let next_stmt = block.stmts.get(index + 1);
            self.lower_stmt(stmt, next_stmt, in_branch)?;
        }
        Ok(())
    }

    fn lower_stmt(
        &mut self,
        stmt: &Stmt,
        next_stmt: Option<&Stmt>,
        in_branch: bool,
    ) -> Result<(), String> {
        match stmt {
            Stmt::Local(local) => self.lower_local(local, next_stmt, in_branch),
            Stmt::Expr(expr, _) => self.lower_expr_stmt(expr, next_stmt, in_branch),
            Stmt::Item(_) => Err("不支持在函数体内定义子项".to_string()),
            Stmt::Macro(_) => Err("暂不支持宏语句".to_string()),
        }
    }

    fn lower_local(
        &mut self,
        local: &syn::Local,
        next_stmt: Option<&Stmt>,
        _in_branch: bool,
    ) -> Result<(), String> {
        let Pat::Ident(PatIdent { ident, .. }) = &local.pat else {
            return Err("仅支持标识符形式的变量绑定".to_string());
        };

        let init = match local.init.as_ref() {
            Some(init) => init.expr.as_ref(),
            None => return Err("仅支持带初始化的 let 语句".to_string()),
        };

        let should_target_return =
            next_stmt.map_or(false, |stmt| is_return_of_ident(stmt, &ident.to_string()));

        let target = if should_target_return { Some(0) } else { None };
        let value_reg = self.lower_expr_with_target(init, target)?;
        self.variables.insert(ident.to_string(), value_reg);

        Ok(())
    }

    fn lower_expr_stmt(
        &mut self,
        expr: &Expr,
        _next_stmt: Option<&Stmt>,
        in_branch: bool,
    ) -> Result<(), String> {
        match expr {
            Expr::Return(ret) => self.lower_return(ret, in_branch),
            Expr::If(expr_if) => self.lower_if(expr_if, in_branch),
            Expr::Block(expr_block) => self.lower_block(&expr_block.block, in_branch),
            _ => Err("暂不支持该语句形式".to_string()),
        }
    }

    fn lower_return(&mut self, ret: &ExprReturn, in_branch: bool) -> Result<(), String> {
        match self.ret_ty {
            MirType::I32 => {
                let expr = ret
                    .expr
                    .as_ref()
                    .ok_or_else(|| "返回类型为 i32 时必须返回表达式".to_string())?;
                self.lower_expr_with_target(expr, Some(0))?;
            }
            MirType::Unit => {
                if ret.expr.is_some() {
                    return Err("Unit 返回类型的函数不应返回值".to_string());
                }
            }
        }

        if in_branch {
            let target = self.ensure_return_block();
            if self.current_block != target {
                self.push_statement(self.current_block, Statement::Goto { target });
            }
        } else {
            if let Some(target) = self.return_block {
                if self.current_block != target {
                    self.push_statement(self.current_block, Statement::Goto { target });
                    return Ok(());
                }
            }
            self.push_statement(self.current_block, Statement::Return);
        }
        Ok(())
    }

    fn lower_if(&mut self, expr_if: &ExprIf, _in_branch: bool) -> Result<(), String> {
        let cond_reg = self.lower_expr_with_target(&expr_if.cond, None)?;
        let true_block = self.mir.add_block();
        let false_block = self.mir.add_block();
        let merge_block = self.mir.add_block();

        self.push_statement(
            self.current_block,
            Statement::SwitchInt {
                cond: Operand::Register(cond_reg),
                true_target: true_block,
                false_target: false_block,
            },
        );

        let previous_return_block = self.return_block;
        if self.return_block.is_none() {
            self.return_block = Some(merge_block);
        }

        self.current_block = true_block;
        self.lower_block(&expr_if.then_branch, true)?;
        self.ensure_block_ends_with_goto(true_block, merge_block);

        self.current_block = false_block;
        if let Some((_, else_expr)) = &expr_if.else_branch {
            match else_expr.as_ref() {
                Expr::Block(block) => self.lower_block(&block.block, true)?,
                Expr::If(nested_if) => self.lower_if(nested_if, true)?,
                _ => return Err("else 分支仅支持块或 if 表达式".to_string()),
            }
        }
        self.ensure_block_ends_with_goto(false_block, merge_block);

        self.current_block = merge_block;
        if previous_return_block.is_some() {
            self.return_block = previous_return_block;
        }
        debug_assert!(self.current_block == merge_block);

        Ok(())
    }

    fn ensure_block_ends_with_goto(&mut self, block_idx: usize, target: usize) {
        if let Some(block) = self.mir.blocks.get(block_idx) {
            if matches!(
                block.statements.last(),
                Some(Statement::Goto { .. }) | Some(Statement::Return)
            ) {
                return;
            }
            if matches!(block.statements.last(), Some(Statement::SwitchInt { .. })) {
                return;
            }
        }
        self.push_statement(block_idx, Statement::Goto { target });
    }

    fn ensure_return_block(&mut self) -> usize {
        if let Some(idx) = self.return_block {
            idx
        } else {
            let idx = self.mir.add_block();
            self.return_block = Some(idx);
            idx
        }
    }

    fn lower_expr_with_target(&mut self, expr: &Expr, target: Option<u32>) -> Result<u32, String> {
        match expr {
            Expr::Binary(binary) => self.lower_binary(binary, target),
            Expr::Paren(ExprParen { expr, .. }) => self.lower_expr_with_target(expr, target),
            Expr::Path(path) => self.lower_path(path, target),
            Expr::Lit(lit) => self.lower_literal(lit, target),
            _ => Err("暂不支持的表达式形式".to_string()),
        }
    }

    fn lower_binary(&mut self, binary: &ExprBinary, target: Option<u32>) -> Result<u32, String> {
        let left_op = self.lower_operand(&binary.left)?;
        let right_op = self.lower_operand(&binary.right)?;

        match &binary.op {
            BinOp::Add(_) => self.emit_binary(BinOpKind::Add, left_op, right_op, target),
            BinOp::Sub(_) => self.emit_binary(BinOpKind::Sub, left_op, right_op, target),
            BinOp::Mul(_) => self.emit_binary(BinOpKind::Mul, left_op, right_op, target),
            BinOp::Gt(_) => self.emit_compare(CmpOpKind::Gt, left_op, right_op, target),
            BinOp::Lt(_) => self.emit_compare(CmpOpKind::Lt, left_op, right_op, target),
            BinOp::Ge(_) => self.emit_compare(CmpOpKind::Ge, left_op, right_op, target),
            BinOp::Le(_) => self.emit_compare(CmpOpKind::Le, left_op, right_op, target),
            BinOp::Eq(_) => self.emit_compare(CmpOpKind::Eq, left_op, right_op, target),
            _ => Err("暂不支持该二元运算".to_string()),
        }
    }

    fn emit_binary(
        &mut self,
        op: BinOpKind,
        left: Operand,
        right: Operand,
        target: Option<u32>,
    ) -> Result<u32, String> {
        let dest = target.unwrap_or_else(|| self.reg_alloc.alloc());
        self.push_statement(
            self.current_block,
            Statement::Assign {
                dest,
                value: RValue::Binary { op, left, right },
            },
        );
        Ok(dest)
    }

    fn emit_compare(
        &mut self,
        op: CmpOpKind,
        left: Operand,
        right: Operand,
        target: Option<u32>,
    ) -> Result<u32, String> {
        let dest = target.unwrap_or_else(|| self.reg_alloc.alloc());
        self.push_statement(
            self.current_block,
            Statement::Assign {
                dest,
                value: RValue::Compare { op, left, right },
            },
        );
        Ok(dest)
    }

    fn lower_operand(&mut self, expr: &Expr) -> Result<Operand, String> {
        match expr {
            Expr::Lit(lit) => {
                let value = parse_literal(lit)?;
                Ok(Operand::Immediate(value))
            }
            Expr::Path(path) => {
                let reg = self.resolve_path(path)?;
                Ok(Operand::Register(reg))
            }
            Expr::Paren(ExprParen { expr, .. }) => self.lower_operand(expr),
            _ => {
                let reg = self.lower_expr_with_target(expr, None)?;
                Ok(Operand::Register(reg))
            }
        }
    }

    fn lower_path(&mut self, path: &ExprPath, target: Option<u32>) -> Result<u32, String> {
        let reg = self.resolve_path(path)?;
        if let Some(dest) = target {
            if dest == reg {
                return Ok(dest);
            }
            self.push_statement(
                self.current_block,
                Statement::Assign {
                    dest,
                    value: RValue::Move(Operand::Register(reg)),
                },
            );
            Ok(dest)
        } else {
            Ok(reg)
        }
    }

    fn lower_literal(&mut self, lit: &ExprLit, target: Option<u32>) -> Result<u32, String> {
        let value = parse_literal(lit)?;
        let dest = target.unwrap_or_else(|| self.reg_alloc.alloc());
        self.push_statement(
            self.current_block,
            Statement::Assign {
                dest,
                value: RValue::Move(Operand::Immediate(value)),
            },
        );
        Ok(dest)
    }

    fn resolve_path(&self, path: &ExprPath) -> Result<u32, String> {
        if path.qself.is_some() {
            return Err("不支持限定路径".to_string());
        }
        let ident = path
            .path
            .get_ident()
            .ok_or_else(|| "暂不支持复杂路径".to_string())?
            .to_string();
        self.variables
            .get(&ident)
            .copied()
            .ok_or_else(|| format!("未找到变量 `{ident}`"))
    }

    fn push_statement(&mut self, block_idx: usize, stmt: Statement) {
        if let Some(block) = self.mir.blocks.get_mut(block_idx) {
            block.statements.push(stmt);
        }
    }
}

fn map_type(ty: &Type) -> Result<MirType, String> {
    match ty {
        Type::Path(type_path) => {
            if let Some(ident) = type_path.path.get_ident() {
                match ident.to_string().as_str() {
                    "i32" => Ok(MirType::I32),
                    "()" => Ok(MirType::Unit),
                    _ => Err(format!("暂不支持类型 `{}`", ident)),
                }
            } else {
                Err("暂不支持复合类型".to_string())
            }
        }
        Type::Tuple(tuple) if tuple.elems.is_empty() => Ok(MirType::Unit),
        _ => Err("暂不支持该类型".to_string()),
    }
}

fn parse_literal(lit: &ExprLit) -> Result<i32, String> {
    match &lit.lit {
        Lit::Int(int_lit) => int_lit
            .base10_parse::<i32>()
            .map_err(|_| "整数文字超出 i32 范围".to_string()),
        _ => Err("暂不支持该字面量".to_string()),
    }
}

fn is_return_of_ident(stmt: &Stmt, ident: &str) -> bool {
    match stmt {
        Stmt::Expr(Expr::Return(ret), _) => match ret.expr.as_deref() {
            Some(Expr::Path(path)) => path
                .path
                .get_ident()
                .map(|name| name == ident)
                .unwrap_or(false),
            _ => false,
        },
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use super::generate_mir;

    fn expect_output(input: &str, expected: &str) {
        let source = std::fs::read_to_string(input)
            .unwrap_or_else(|err| panic!("读取输入文件失败 {}: {}", input, err));
        let expected_output = std::fs::read_to_string(expected)
            .unwrap_or_else(|err| panic!("读取期望输出失败 {}: {}", expected, err));
        let generated = generate_mir(&source).expect("生成 MIR 失败");
        assert_eq!(generated, expected_output);
    }

    #[test]
    fn lab4_example1() {
        expect_output(
            "tests/lab4-tests/example1.rs",
            "tests/lab4-tests/example1.expected",
        );
    }

    #[test]
    fn lab4_example4() {
        expect_output(
            "tests/lab4-tests/example4.rs",
            "tests/lab4-tests/example4.expected",
        );
    }

    #[test]
    fn lab4_example10() {
        expect_output(
            "tests/lab4-tests/example10.rs",
            "tests/lab4-tests/example10.expected",
        );
    }
}
